#include <iostream>
#include <cmath>
#include <cuComplex.h>

// FFT kernel
__global__ void fft_kernel(cuComplex* data, int N) {
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < N) {
        cuComplex temp = make_cuComplex(0.0f, 0.0f);

        for (int k = 0; k < N; k++) {
            float realPart = cosf(-2.0f * M_PI * idx * k / N);
            float imagPart = sinf(-2.0f * M_PI * idx * k / N);
            cuComplex twiddle = make_cuComplex(realPart, imagPart);

            temp = cuCaddf(temp, cuCmulf(data[k], twiddle));
        }

        data[idx] = temp;
    }
}

// 比较三个结果并投票表决
cuComplex vote_result(cuComplex a, cuComplex b, cuComplex c) {
    // 如果三个结果中有两个相同，则返回相同的那个
    if (cuCrealf(a) == cuCrealf(b) || cuCrealf(a) == cuCrealf(c)) {
        return a;
    }
    return b; // 如果a与b、c都不相同，则返回b(假设b和c相同)
}

int main() {
    const int N = 8;
    cuComplex* h_data_1 = new cuComplex[N];
    cuComplex* h_data_2 = new cuComplex[N];
    cuComplex* h_data_3 = new cuComplex[N];
    cuComplex* h_final_result = new cuComplex[N];

    // 初始化数据
    for (int i = 0; i < N; i++) {
        h_data_1[i] = make_cuComplex(1.0f, 0.0f);
        h_data_2[i] = make_cuComplex(1.0f, 0.0f);
        h_data_3[i] = make_cuComplex(1.0f, 0.0f);
    }

    cuComplex* d_data_1;
    cuComplex* d_data_2;
    cuComplex* d_data_3;
    cudaMalloc((void**)&d_data_1, N * sizeof(cuComplex));
    cudaMalloc((void**)&d_data_2, N * sizeof(cuComplex));
    cudaMalloc((void**)&d_data_3, N * sizeof(cuComplex));
    
    // 三模冗余：将相同数据拷贝到三个不同的设备内存区域
    cudaMemcpy(d_data_1, h_data_1, N * sizeof(cuComplex), cudaMemcpyHostToDevice);
    cudaMemcpy(d_data_2, h_data_2, N * sizeof(cuComplex), cudaMemcpyHostToDevice);
    cudaMemcpy(d_data_3, h_data_3, N * sizeof(cuComplex), cudaMemcpyHostToDevice);

    // 三模冗余：运行三个相同的FFT计算
    fft_kernel<<<1, N>>>(d_data_1, N);
    fft_kernel<<<1, N>>>(d_data_2, N);
    fft_kernel<<<1, N>>>(d_data_3, N);

    // 将结果拷贝回主机
    cudaMemcpy(h_data_1, d_data_1, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_data_2, d_data_2, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_data_3, d_data_3, N * sizeof(cuComplex), cudaMemcpyDeviceToHost);

    // 三模冗余：对三个结果进行投票表决
    for (int i = 0; i < N; i++) {
        h_final_result[i] = vote_result(h_data_1[i], h_data_2[i], h_data_3[i]);
        std::cout << "Final Result[" << i << "] = " 
                  << cuCrealf(h_final_result[i]) << " + " 
                  << cuCimagf(h_final_result[i]) << "i" << std::endl;
    }

    // 检查是否有错误被纠正
    bool error_detected = false;
    for (int i = 0; i < N; i++) {
        if (!(cuCrealf(h_data_1[i]) == cuCrealf(h_data_2[i]) && 
              cuCimagf(h_data_1[i]) == cuCimagf(h_data_2[i]) &&
              cuCrealf(h_data_1[i]) == cuCrealf(h_data_3[i]) &&
              cuCimagf(h_data_1[i]) == cuCimagf(h_data_3[i]))) {
            error_detected = true;
            break;
        }
    }
    
    if (error_detected) {
        std::cout << "Warning: Errors detected and corrected by TMR!" << std::endl;
    } else {
        std::cout << "All three modules produced identical results." << std::endl;
    }

    // 清理
    delete[] h_data_1;
    delete[] h_data_2;
    delete[] h_data_3;
    delete[] h_final_result;
    cudaFree(d_data_1);
    cudaFree(d_data_2);
    cudaFree(d_data_3);

    return 0;
}